package cn.seecoder.ai.algorithm.ml

import cn.seecoder.ai.enums.TrainParamEnum
import cn.seecoder.ai.model.bo.TrainParamsBO
import cn.seecoder.ai.utils.TrainParamReader.readParam
import org.apache.spark.ml.evaluation.{BinaryClassificationEvaluator, MulticlassClassificationEvaluator}
import org.apache.spark.ml.{Pipeline, PipelineModel, PipelineStage}
import org.apache.spark.ml.feature.VectorIndexer
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
import org.slf4j.{Logger, LoggerFactory}

import java.util
import org.apache.spark.ml.linalg.Vector
/**
 * todo: 实现分类模型的train方法，使用了模板设计模式
 *
 * @author   fanyanpeng
 * @date 2023/4/17 5:12
 * @param null
 * @return
 */
class BaseClassification extends BaseAlgorithm {

  val log: Logger = LoggerFactory.getLogger(getClass)


  //获取标签种类数目
  private def getLabelUniqueNum (dataFrame: DataFrame): Int = {
    dataFrame.select("label").distinct().show()
    dataFrame.select("label").distinct().count().toInt
  }

  // 获取特征维度
  private def getFeatureDimensionNum(dataFrame: DataFrame): Int = {
    dataFrame.select("features").first().getAs[Vector](0).size
  }

  /**
   * 切分数据集，前后部分切分
   * @param dataFrame
   * @param trainingDataSetOccupy
   * @return
   */
  private def splitDataSet(dataFrame: DataFrame, trainingDataSetOccupy: Double): Array[Dataset[Row]] = {
    val count = dataFrame.count()
    val subTrainSetCount= (count * trainingDataSetOccupy).toInt
    var subTrainSet = dataFrame.limit(subTrainSetCount)
    var subTestSet = dataFrame.except(subTrainSet)
    Array(subTrainSet, subTestSet)
  }

  override def train(sparkSession: SparkSession,
            libSvmFileUri: String,
            trainParams: TrainParamsBO): (PipelineModel,util.Map[String, Object]) ={

    val trainSet = sparkSession.read.format("libsvm").load(libSvmFileUri)

    //获取最大种类，若读取失败，获取默认值
    val maxCategories = trainParams.maxCategories


    val featureIndexer = new VectorIndexer()
      .setInputCol("features")
      .setOutputCol("indexedFeatures")
      .setMaxCategories(maxCategories).setHandleInvalid("keep")

    trainSet.show()
    trainParams.labelUniqueNum = getLabelUniqueNum(trainSet)

    //读取训练集占比
    val trainingDataSetOccupy: Double = trainParams.trainingDataSetOccupy

    //训练集切分，修改为前后部分切分
    var Array(subTrainSet: Dataset[Row], subTestSet: Dataset[Row]) = splitDataSet(trainSet, trainingDataSetOccupy)

    subTrainSet.show(10)
    subTestSet.show(10)


    //本地添加参数
    trainParams.featureDimensionNum = getFeatureDimensionNum(subTrainSet)


    subTrainSet.select("label").distinct().show()
    subTestSet.select("label").distinct().show()

    /** **************************************** */
    //需要由子类实现
    val machineLearningStage = buildMachineLearningStage(trainParams)

    /** **************************************** */

    val pipeline = new Pipeline()
      .setStages(Array(featureIndexer, machineLearningStage))

    // Train model. This also runs the indexers.
    val pipelineModel: PipelineModel = pipeline.fit(subTrainSet)

    val resultMap: util.Map[String, Object] = evaluate(pipelineModel, subTestSet, machineLearningStage, trainParams)

    return (pipelineModel, resultMap)

  }


  def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {
    //由子类实现
    null
  }

  def evaluate(pipelineModel: PipelineModel, subTestSet: DataFrame, machineLearningStage: PipelineStage, trainParams: TrainParamsBO): util.Map[String, Object] ={

    val resultMap: util.Map[String, Object] = new util.HashMap[String, Object]()
    // Make predictions.
    val predictions = pipelineModel.transform(subTestSet)

    predictions.show(10)

    // 若只有两类label，这是一个二分类问题，使用二分类评价器
    if(trainParams.labelUniqueNum == 2){
      val evaluator = new BinaryClassificationEvaluator().setLabelCol("label").setRawPredictionCol("rawPrediction")

      val areaUnderROC = evaluator.setMetricName("areaUnderROC").evaluate(predictions)
      val areaUnderPR = evaluator.setMetricName("areaUnderPR").evaluate(predictions)


      resultMap.put("areaUnderROC", areaUnderROC.toString)
      resultMap.put("areaUnderPR", areaUnderPR.toString)
    }
    else{

      val evaluator = new MulticlassClassificationEvaluator()
        .setLabelCol("label")
        .setPredictionCol("prediction")

      val weightedRecall = evaluator.setMetricName("weightedRecall").evaluate(predictions);
      val accuracy = evaluator.setMetricName("accuracy").evaluate(predictions)
      log.info("Test Error = " + (1.0 - accuracy))
      val weightedPrecision = evaluator.setMetricName("weightedPrecision").evaluate(predictions);
      val f1 = evaluator.setMetricName("f1").evaluate(predictions);

      val model = pipelineModel.stages(1).asInstanceOf[machineLearningStage.type]
      log.info("\n\nLearned classification model:\n" + model.toString());


      resultMap.put("Accuracy", accuracy.toString)
      resultMap.put("weightedPrecision", weightedPrecision.toString)
      resultMap.put("weightedRecall", weightedRecall.toString)
      resultMap.put("f1", f1.toString)
    }

    resultMap
  }

}
